

from .config_manager import ConfigManager
import asyncio  # 用于异步编程
import os
import paramiko  # SSH 协议库，用于远程执行命令
import logging
import threading  # 用于线程控制
import time  # 处理时间与定时
from typing import Optional, Callable, AsyncGenerator  # 类型提示
from dataclasses import dataclass  # 简化数据类的声明
from enum import Enum  # 枚举支持
import traceback  # 用于异常回溯
import re

from .config_manager import Task, Host
from .auto_reply import ReplyGenerator
from .expected_match import ExpectedMatch

logger = logging.getLogger(__name__)  # 获取日志记录器实例

@dataclass
class Stream:
    """用于存放流式命令执行信息的数据结构"""
    command: str   # 记录命令内容
    host_name: str # 记录主机名
    task_name: str = ''  # 记录任务名称
    is_running: bool = True  # 是否仍在运行

@dataclass
class StreamingResult:
    """用于存放流式命令执行信息的数据结构"""
    command: str   # 记录命令内容
    host_name: str # 记录主机名
    task_name: str = None  # 记录任务名称
    is_running: bool = True  # 是否仍在运行
    output_lines: list = None  # 存放命令标准输出的行
    error_lines: list = None   # 存放命令标准错误的行
    start_time: float = None   # 命令开始时间
    last_output_time: float = None  # 最近有输出的时间
    write_to_file: bool = True  # 是否写入文件
    show_realtime: bool = False  # 是否实时打印输出
    stdout_file: str = None  # 记录输出的文件路径
    stderr_file: str = None  # 记录错误的文件路径
    line_timeout: int = 10  # 无新输出的超时时间（秒）
    max_duration: int = 60  # 最大运行时间（秒）
    kill_timeout: int = 1  # 超时kill后的超时时间，用于捕获kill后的输出
    cmd_output_dir: str = None  # 程序输出路径


    def __post_init__(self):
        """初始化时填充默认值，避免空指针问题"""
        if self.output_lines is None:
            self.output_lines = []
        if self.error_lines is None:
            self.error_lines = []
        if self.start_time is None:
            self.start_time = time.time()  # 默认为当前时间
        self.last_output_time = self.start_time  # 初始化最近输出时间
        if self.write_to_file:
            if self.cmd_output_dir is None:
                self.cmd_output_dir = f"cmd_output_log"
            os.makedirs(self.cmd_output_dir, exist_ok=True)

            self.stderr_file = f"{self.cmd_output_dir}/{self.host_name}__{self.task_name}__stderr__{int(self.start_time)}.log"  # 生成文件名    
            self.stdout_file = f"{self.cmd_output_dir}/{self.host_name}__{self.task_name}__stdout__{int(self.start_time)}.log"  # 生成文件名

class streamCommandExecutor:
    """增强版命令执行器，支持异步和流式命令实时输出处理"""

    def __init__(self, ssh_manager):
        self.ssh_manager = ssh_manager  # 管理主机SSH连接的管理器对象
        self.active_streams = dict()  # 用于追踪活跃的流式命令，通过 {host:task} 标识
        self._stop_event = threading.Event()  # 用于全局中断
        # self._lock = threading.Lock()  # 用于线程安全操作

    async def execute_streaming_command(
        self, 
        host: Host, 
        task: Task,
        cmd_output_dir: str = '',
        output_callback: Optional[Callable] = None,
    ) -> StreamingResult:
        """
        执行持续输出的命令（如 tail -f），以流式方式异步处理输出。
        
        Args:
            host: 主机连接配置对象
            task: 任务配置对象
            command: 待执行的 shell 命令字符串
            max_duration: 运行最大时间限制
            show_realtime: 是否实时展示
            write_to_file: 是否写入文件
            line_timeout: 若多少秒无更新则判断为“异常挂死”
        Returns:
            StreamingResult: 记录所有流式输出和异常的结果对象
        """
        command = task.command
        streaming_config = getattr(task, 'streaming_config', {})
        shell_destroy = streaming_config.get('shell_destroy', True)
        shell_inherit_task = streaming_config.get('shell_inherit_task', '') # shell来源的的Task名称，为空则新建shell
        max_duration = streaming_config.get('max_duration', 60)
        line_timeout = streaming_config.get('line_timeout', 10)
        kill_timeout = streaming_config.get('kill_timeout', 1)
        show_realtime = streaming_config.get('show_realtime', False)
        write_to_file = streaming_config.get('write_to_file', True)


        stream_key = host.name +"|"+ task.name
  
        self.active_streams[stream_key] = Stream(command=command, host_name=host.name, task_name=task.name)
        
        result = StreamingResult(command=command, 
                                 host_name=host.name, 
                                 task_name=task.name, 
                                 max_duration=max_duration,
                                 line_timeout=line_timeout,
                                 kill_timeout=kill_timeout,
                                 show_realtime=show_realtime,
                                 write_to_file=write_to_file,
                                 cmd_output_dir=cmd_output_dir)

        output_callback = self._realtime_callback() if output_callback is None else output_callback 
        reply_generator = ReplyGenerator(task.interactive_response)
        expected_match = ExpectedMatch(task.expected, task.unexpected)

        shell_key = ""
        try:
            # 1. 获取交互式 shell 会话
            shell = None
            if shell_inherit_task:
                # shell_key = host.name +"|"+ shell_inherit_task
                shell_key= host.name +"|"+ host.host+":"+str(host.port)+"|"+shell_inherit_task

                if self.ssh_manager.has_shell(shell_key):
                    shell = self.ssh_manager.get_shell(host,shell_inherit_task)
                    logger.info(f"[{host.name}|{task.name}] 继承shell: [{shell_key}]")
                else:
                    logger.warning(f"未找到继承的shell: [{shell_key}]")
                    shell = self.ssh_manager.get_shell(host,task.name)
                    shell_key = host.name +"|"+ task.name
                    logger.info(f"[{shell_key}] 已创建新shell")
            else:
                shell = self.ssh_manager.get_shell(host,task.name)
                shell_key = host.name +"|"+ task.name
                logger.info(f"[{shell_key}] 已创建新shell")
            
            await asyncio.sleep(0.5)  # 等待shell初始化
            shell.settimeout(1.0)  # 设置短超时，防止阻塞

            logger.info(f"[{shell_key}] 向shell发送命令: {command}")
            # 2. 发送要执行的命令
            shell.send(command + '\n')

            # 3. 进入异步流式读输出
            await self._read_streaming_output(
                shell, host, task, result, 
                output_callback=output_callback,
                reply_generator=reply_generator,
                expected_match=expected_match
            )
            
        except Exception as e:
            # 出现异常时记录到错误行
            logger.error(f"执行命令 {command} 时出错: {str(e)}")
            logger.error(traceback.format_exc())
        finally:
            # result.is_running = False  # 标志为已结束
            self.active_streams[stream_key].is_running = False # 标志为已结束
            # 检查输出是否符合期望
            if expected_match:
                ok,exception = expected_match.check()
                if not ok:
                    result.error_lines.append(f"{exception}")
                    logger.warning(f"[{host.name}|{task.name}] {exception}")

            # 是否销毁shell
            if shell_destroy:
                self.ssh_manager.close_shell(host,task.name)
                logger.info(f"[{shell_key}] 已销毁shell")
        
        return result

    async def _read_streaming_output(
        self, 
        shell, 
        host: Host,
        task: Task,
        result: StreamingResult, 
        output_callback: Optional[Callable],
        reply_generator: Optional[ReplyGenerator],
        expected_match: Optional[ExpectedMatch],
    ):
        """流式读取命令输出的主循环和控制逻辑"""
        start_time = time.time()
        last_output_time = start_time

        # while result.is_running:
        stream_key = host.name +"|"+ task.name
        try:
            while self.active_streams[stream_key].is_running:
                current_time = time.time()
                
                data = ''
                # 超时逻辑1: 总执行时长是否超标
                if (current_time - start_time) > result.max_duration:
                    result.output_lines.append(f"[TIMEOUT] 达到最大执行时间 {result.max_duration}s,自动中断")
                    logger.warning(f"[{host.name}|{task.name}] 达到最大执行时间 {result.max_duration}s,自动中断")
                    shell.send("\x03")
                    await asyncio.sleep(result.kill_timeout)
                    # 读取剩余数据
                    data = shell.recv(40960).decode('utf-8', errors='ignore') 
                    if expected_match:
                        expected_match.match(data)       
                    last_output_time = current_time
                    result.last_output_time = current_time
                    if output_callback:
                        await self._safe_callback(output_callback,"stdout", data, result)
                    # result.is_running = False
                    break

                # 超时逻辑2: 若(N秒)无新内容输出视为挂起
                if current_time - last_output_time > result.line_timeout:
                    result.output_lines.append(f"[TIMEOUT] 无输出超时 {result.line_timeout}s,自动中断")
                    logger.warning(f"[{host.name}|{task.name}] 无输出超时 {result.line_timeout}s,自动中断")
                    shell.send("\x03")
                    await asyncio.sleep(result.kill_timeout)
                    # 读取剩余数据
                    data = shell.recv(40960).decode('utf-8', errors='ignore') 
                    if expected_match:
                        expected_match.match(data)       
                    last_output_time = current_time
                    result.last_output_time = current_time
                    if output_callback:
                        await self._safe_callback(output_callback,"stdout", data, result)
                    # result.is_running = False
                    break
                
                # 退出逻辑1：所有期望输出均已满足,且每个都匹配到了指定次数(已经成功，提前退出)
                if expected_match:
                    if expected_match.check_expected():
                        logger.info(f"[{host.name}|{task.name}] 所有期望输出均已满足,且每个都匹配到了指定次数(已经成功，提前退出)")
                        break
                
                # 退出逻辑2: 有不期望输出匹配到，且已经超过最大容忍次数(已经失败，提前退出)
                if expected_match:
                    if not expected_match.check_unexpected():
                        logger.warning(f"[{host.name}|{task.name}] 有不期望输出匹配到，且已经超过最大容忍次数(已经失败，提前退出)")
                        break


                # 检查标准输出是否有新数据
                if shell.recv_ready():
                    data = shell.recv(40960).decode('utf-8', errors='ignore')
                    data = data.replace('\r\n', '\n')
                    if data:
                        last_output_time = current_time
                        result.last_output_time = current_time
                        # 流式输出的结果，将不会打印到控制台，而是直接重定向到文件，防止过多数据导致内存溢出。
                        # result.output_lines.append(data) 
                        if output_callback:
                            await self._safe_callback(output_callback,"stdout", data, result)
                        
                        # 自动回复
                        if reply_generator:
                            reply = reply_generator.generate_reply(data)
                            if reply:
                                shell.send(reply + '\n')
                                logger.info(f"[{host.name}|{task.name}] 自动回复: {reply}")
                                await asyncio.sleep(0.2)  # 等待处理
                        
                # 检查是否匹配期望输出
                if expected_match:
                    expected_match.match(data)                            
                        
                # 检查标准错误流
                if shell.recv_stderr_ready():
                    error_data = shell.recv_stderr(40960).decode('utf-8', errors='ignore')
                    if error_data:
                        last_output_time = current_time
                        result.last_output_time = current_time
                        if output_callback:
                            await self._safe_callback(output_callback,"stderr", error_data, result)
                        # 流式输出的结果，将不会打印到控制台，而是直接重定向到文件，防止过多数据导致内存溢出。
                        # result.error_lines.extend(error_data.strip().split('\n'))

                # 短暂休眠避免过高CPU占用
                await asyncio.sleep(0.1)    
                
        except Exception as e:
            # 读取及解码过程中的异常
            # result.error_lines.append(f"读取输出异常: {str(e)}")
            logger.error(f"读取输出异常: {str(e)}")
            logger.error(traceback.format_exc())
            # result.is_running = False
            self.active_streams[stream_key].is_running = False
        finally:
            self.active_streams[stream_key].is_running = False
            
            
                

    async def _safe_callback(self, callback, data_type: str, data: str, result: StreamingResult):
        """安全地执行回调函数，自动兼容同步/异步函数，并捕捉其异常"""
        try:
            if asyncio.iscoroutinefunction(callback):
                await callback(data_type, data, result)  # 异步回调
            else:
                callback(data_type, data, result)        # 同步回调
        except Exception as e:
            # result.error_lines.append(f"回调函数异常: {str(e)}")
            logger.error(f"回调函数异常: {str(e)}")

    def _realtime_callback(self):
        """
        为流式命令创建实时输出回调函数
            
        Returns:
            function: 回调函数
        """
        def remove_ansi_codes(text):
            """
            去除文本中的ANSI颜色控制码。
            """
            ansi_escape = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]')
            return ansi_escape.sub('', text)
    
        async def callback(data_type: str, data: str, result: StreamingResult):
            """实时输出回调"""
            data = remove_ansi_codes(data)
            if data_type=="stdout":
                if result.write_to_file and result.stdout_file:
                    with open(result.stdout_file, "a") as f:
                        f.write(data)
                        f.flush()
                if result.show_realtime:
                    logger.info(f"[{result.host_name}|{result.task_name}] 输出: \n{data}")
                    
            elif data_type=="stderr":
                if result.write_to_file and result.stderr_file:
                    with open(result.stderr_file, "a") as f:
                        f.write(data)
                        f.flush()
                if result.show_realtime:
                    logger.error(f"[{result.host_name}]:[{result.task_name}] 报错| \n{data}")

        return callback

    

    
    def stop_streaming_command(self, host_name: str, task_name: str):
        """停止某个具体的流式命令任务"""
        key = f"{host_name}|{task_name}"
        if key in self.active_streams:
            self.active_streams[key].is_running = False

    def stop_all_streaming(self):
        """停止所有活跃中的流式命令实例"""
        # for result in self.active_streams.values():
        #     result.is_running = False
        for key in self.active_streams.keys():
            self.active_streams[key].is_running = False
            # del self.active_streams[key]
            
        self._stop_event.set()  # 通用线程信号